import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

def assign_cluster(x, c):
    """
    将样本 x 分配到最近的聚类中心
    """
    distances = np.linalg.norm(x[:, np.newaxis] - c, axis=2)  # shape: (n_samples, K)
    y = np.argmin(distances, axis=1)
    return y


def Kmean(data, K, epsilon=1e-4, max_iteration=100):
    """
    K-Means 聚类算法
    """
    n_samples, n_features = data.shape
    np.random.seed(42)
    centers = data[np.random.choice(n_samples, K, replace=False)]

    for i in range(max_iteration):
        # Step 1: 分配簇
        labels = assign_cluster(data, centers)

        # Step 2: 更新中心
        new_centers = np.array([
            data[labels == k].mean(axis=0) if np.any(labels == k) else centers[k]
            for k in range(K)
        ])

        # Step 3: 判断收敛
        shift = np.linalg.norm(new_centers - centers)
        if shift < epsilon:
            print(f"第 {i+1} 次迭代后收敛，中心移动量 {shift:.6f}")
            break

        centers = new_centers

    return centers, labels


if __name__ == "__main__":
    # 生成二维数据
    data, _ = make_blobs(n_samples=200, centers=3, n_features=2, random_state=42)

    centers, labels = Kmean(data, K=3)
    print("最终聚类中心：\n", centers)

    # ======================
    # 可视化
    # ======================
    plt.figure(figsize=(8, 6))

    # 绘制每个簇的样本点
    for k in range(3):
        plt.scatter(data[labels == k, 0], data[labels == k, 1], label=f'Cluster {k}', alpha=0.6)

    # 绘制中心点
    plt.scatter(centers[:, 0], centers[:, 1], c='black', s=200, marker='X', label='Centers')

    plt.title("K-Means Clustering Result")
    plt.xlabel("Feature 1")
    plt.ylabel("Feature 2")
    plt.legend()
    plt.grid(True)
    plt.show()
